import sys

#sys.path.insert(0, "build/lib.macosx-10.4-i386-2.5")
#sys.path.insert(0, "../../../jj_coord/build/lib.macosx-10.4-i386-2.5")
sys.path.insert(0, "../build/lib.linux-i686-2.5/")

import coords


import pywcs
import pyfits
import numpy as np

import pywcsgrid


import pylab

def tolatex(s):
    return [s1.replace(r"\u","").replace(r"\d","") for s1 in s]


def plot_grid(ax, g, draw_grid=True):

    if draw_grid:
        grid_lines1 = [ ax.plot(x, y, "k-") for x, y in g.axis1["grid_lines"]]
        grid_lines2 = [ ax.plot(x, y, "k-") for x, y in g.axis2["grid_lines"]]
        grid_lines = grid_lines1 + grid_lines2
    else:
        grid_lines = []


    ax.set_xticks(g.axis2["ticks_bottom"])
    ax.set_xticklabels(tolatex(g.axis2["ticklabels_bottom"]))

    ax.set_yticks(g.axis1["ticks_left"])
    ax.set_yticklabels(tolatex(g.axis1["ticklabels_left"]))


    ax.xaxis.tick_bottom()
    ax.yaxis.tick_left()
    
    return grid_lines



def plot_fits2(fname):
    f = pyfits.open(fname)
    d, h = f[0].data, f[0].header

    d = d.astype("d")
    d[np.isnan(d)] = -1.

    w = pywcs.WCS(h)

    fig1 = pylab.figure(2, figsize=(12,4))
    fig1.subplots_adjust(wspace=0.3)

    ax1 = pylab.subplot(1,3,1)
    ax2 = pylab.subplot(1,3,2)
    ax3 = pylab.subplot(1,3,3)

    ax1.imshow(d, origin="lower", interpolation="nearest")
    ax2.imshow(d, origin="lower", interpolation="nearest")
    ax3.imshow(d, origin="lower", interpolation="nearest")

    g1 = pywcsgrid.grid_gal(w, 0.5, 0.5, h["NAXIS1"], h["NAXIS2"],
                            #label_control=(2,1,2,1)
                            )

    plot_grid(ax1, g1)
    ax1.set_xlim(-0.5, h["NAXIS1"]-0.5)
    ax1.set_ylim(-0.5, h["NAXIS2"]-0.5)

    ax1.set_xlabel("Dec.")
    ax1.set_ylabel("R.A.")

    g2 = pywcsgrid.grid_gal(w, 0.5, 0.5, h["NAXIS1"], h["NAXIS2"],
                            #label_control=(2,1,2,1),
                            label_density=(12, 6),
                            coord_format = "AA",
                            grid1 = [270, 285, 300],
                            grid2 = [-70, -65],
                            )


    plot_grid(ax2, g2)
    ax2.set_xlim(-0.5, h["NAXIS1"]-0.5)
    ax2.set_ylim(-0.5, h["NAXIS2"]-0.5)
    ax2.set_xlabel("Dec.")

    plot_grid(ax3, g1, draw_grid=False)
    ax3.set_xlabel("Dec.")
    ax3.set_xlim(-0.5, h["NAXIS1"]-0.5)
    ax3.set_ylim(-0.5, h["NAXIS2"]-0.5)

    pylab.draw()


def plot_fits(fname):
    f = pyfits.open(fname)
    d, h = f[0].data, f[0].header

    d = d.astype("d")
    d[np.isnan(d)] = -1.

    w = pywcs.WCS(h)

    fig1 = pylab.figure(1, figsize=(10,6))
    fig1.subplots_adjust(left=0.2)

    ax1 = pylab.subplot(1,2,1)

    ax1.imshow(d, origin="lower", interpolation="nearest")

    g1 = pywcsgrid.grid(w, 0.5, 0.5, h["NAXIS1"], h["NAXIS2"],
                        #label_control=(2,1,2,1)
                        )

    plot_grid(ax1, g1)
    ax1.set_xlim(-0.5, h["NAXIS1"]-0.5)
    ax1.set_ylim(-0.5, h["NAXIS2"]-0.5)

    ax1.set_xlabel("Dec.")
    ax1.set_ylabel("R.A.")


    ax1 = pylab.subplot(1,2,2)

    ax1.imshow(d, origin="lower", interpolation="nearest")

    g1 = pywcsgrid.grid_with_coord_conv(w, 0.5, 0.5, 0.5+h["NAXIS1"], 0.5+h["NAXIS2"],
                                        src_system=coords.FK5,
                                        dest_system=coords.GAL,
                                        )

    plot_grid(ax1, g1)
    ax1.set_xlim(-0.5, h["NAXIS1"]-0.5)
    ax1.set_ylim(-0.5, h["NAXIS2"]-0.5)

    ax1.set_xlabel("Dec.")
    ax1.set_ylabel("R.A.")


    pylab.draw()



if __name__ == "__main__":

    import sys
    fname = sys.argv[1]

    plot_fits(fname)
    #plot_fits2(fname)
    
    pylab.show()
